from .base_laysers import Encoder, Q_Net
import torch.nn as nn
import torch
import torch.nn.functional as F


# #############################  CommNet  ########################################
# CommNet 平均模块
class AvgModel(nn.Module):
    def __init__(self, n_node, din, hidden_dim, dout):
        super(AvgModel, self).__init__()
        self.fc1 = nn.Linear(din, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, dout)

    def forward(self, x, mask):
        att = F.softmax(mask - 9e15 * (1 - mask), dim=2)
        x = F.relu(self.fc1(x))
        out = torch.bmm(att, x)
        out = F.relu(self.fc2(out))

        return out


class CommNet(nn.Module):

    def __init__(self, n_agent, num_inputs, hidden_dim, num_actions):
        super(CommNet, self).__init__()
        self.encoder = Encoder(num_inputs, hidden_dim)
        self.avg = AvgModel(n_agent, hidden_dim, hidden_dim, hidden_dim)
        self.q_net = Q_Net(hidden_dim, num_actions)

    def forward(self, x, mask):
        encoded = self.encoder(x)
        avg = self.avg(encoded, mask)
        q = self.q_net(avg)

        return q, mask